# -*- coding: utf-8 -*-
# ----------------------------------------------------------------
# 文件编码声明：确保中文注释和路径中的中文字符能被正确识别和处理
# ----------------------------------------------------------------

import os
import glob
import warnings
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import mne

# 假设 data_process 模块和 get_subject_id 函数已存在且可用
# 注意：由于移除了 CSV 模式，process_eeg_data 不再被使用，但 get_subject_id 仍用于从文件名提取被试ID。
# from data_process import process_eeg_data, get_subject_id 
# 为了代码的独立性，如果 get_subject_id 依赖于 data_process，可能需要确保它仍在 PATH 中或直接在下方定义。
# 这里假设 get_subject_id 是一个独立函数，且功能简单，或者您已将其路径配置好。
# 如果 data_process 不再存在，请确保 get_subject_id 的实现已包含在某处。

# 假设 get_subject_id 实现了以下功能:
def get_subject_id(filename):
    """从 GDF 文件名 (如 'A01T.gdf') 中提取被试 ID (如 1)。"""
    try:
        # 匹配文件名中的数字部分，例如 A01T.gdf -> 01
        import re
        match = re.search(r'A(\d+)[TE]\.gdf', filename, re.IGNORECASE)
        if match:
            return int(match.group(1))
    except ImportError:
        pass # 如果没有 re 库（不太可能），则返回 None
    
    # 简单的回退逻辑
    return None


# 确保中文正常显示
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']
matplotlib.rcParams['axes.unicode_minus'] = False


def read_concat_raws(gdf_paths):
    """读取并合并列表中的所有 GDF 文件"""
    if len(gdf_paths) == 0:
        raise FileNotFoundError('未找到任何.gdf文件')
    raws = []
    for p in gdf_paths:
        with warnings.catch_warnings():
            # 忽略 MNE 在 GDF 文件中可能发出的 highpass/lowpass 警告
            warnings.filterwarnings(
                'ignore',
                message=r"Highpass cutoff frequency .* greater than lowpass cutoff frequency .*",
                category=RuntimeWarning,
            )
            # 忽略通道名称重复警告
            warnings.filterwarnings(
                'ignore',
                message=r"Channel names are not unique.*",
                category=RuntimeWarning,
            )
            r = mne.io.read_raw_gdf(p, preload=True, verbose=False)
        # 修复可能的截止频率顺序错误
        hp, lp = r.info.get('highpass', None), r.info.get('lowpass', None)
        if hp is not None and lp is not None and hp > lp:
            r.info['highpass'], r.info['lowpass'] = lp, hp
        raws.append(r)
    return mne.io.concatenate_raws(raws)


def get_events_from_annotations(raw):
    """
    从 GDF 文件的标注提取 MNE Events 数组。
    仅保留左手(769)和右手(770)想象试次。
    """
    events, event_id = mne.events_from_annotations(raw, verbose=False)
    
    # 仅保留左(769)或右(770)想象试次
    mask = (events[:, 2] == 769) | (events[:, 2] == 770)
    events = events[mask]
    
    if len(events) == 0:
        raise ValueError('未找到有效的左手或右手想象事件')
    
    return events


def resolve_c3_c4_names(ch_names):
    """在通道名列表中寻找 C3 和 C4 的实际名称（例如：C3. 或 C3_REF）。"""
    lower_to_orig = {c.lower(): c for c in ch_names}
    def find_one(target):
        if target in lower_to_orig:
            return lower_to_orig[target]
        cand = [c for c in ch_names if target in c.lower()]
        return cand[0] if len(cand) > 0 else None
    c3 = find_one('c3')
    c4 = find_one('c4')
    return c3, c4


def band_envelope(x, sfreq=250, win_ms=100):
    """简单能量包络：平方后做移动平均（100ms 窗口），返回同长度。"""
    power = x ** 2
    win = max(1, int(sfreq * win_ms / 1000))
    kernel = np.ones(win) / win
    # 使用 np.convolve 进行移动平均
    return np.convolve(power, kernel, mode='same')


def plot_subject_c3c4(base_dir, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    
    # 仅使用训练 GDF 文件 (*T.gdf)，因为测试文件无类标签
    train_gdf = sorted(glob.glob(os.path.join(base_dir, '*T.gdf')))
    if len(train_gdf) == 0:
        raise FileNotFoundError('未检测到训练 GDF 文件 (*.gdf)，请检查路径是否正确。')

    print('✅ 检测到训练GDF文件，使用 8-13Hz Mu 节律 和 0-9 秒完整时间窗绘图（仅训练数据）')
    
    LOW_CUTOFF, HIGH_CUTOFF = 8, 13 # Mu 节律
    
    # 提取被试 ID
    subject_ids = sorted(set(get_subject_id(os.path.basename(p)) for p in train_gdf if get_subject_id(os.path.basename(p)) is not None))
    
    for sid in subject_ids:
        tg = [p for p in train_gdf if get_subject_id(os.path.basename(p)) == sid]
        if len(tg) == 0: continue
        
        try:
            raw = read_concat_raws(tg)
        except Exception as e:
            print(f"被试 {sid} 读取 GDF 错误: {e}")
            continue

        # 从标注提取事件，无需外部标签文件
        try:
            events = get_events_from_annotations(raw)
        except ValueError as e:
            print(f'被试{sid}: {e}，跳过')
            continue
        
        # 仅取 C3/C4 两通道
        c3_name, c4_name = resolve_c3_c4_names(raw.info['ch_names'])
        if c3_name is None or c4_name is None:
            print(f'被试{sid}: 未找到 C3/C4，跳过')
            continue
        
        raw_pick = raw.copy().pick_channels([c3_name, c4_name], ordered=True)
        # 设置平均参考，并应用 Mu 节律滤波
        raw_pick.set_eeg_reference('average', projection=True, verbose=False)
        raw_pick.filter(LOW_CUTOFF, HIGH_CUTOFF, fir_design='firwin', verbose=False)
        
        # 提取 Epochs (0.0s 到 9.0s 完整时间窗，无基线校正)
        epochs = mne.Epochs(raw_pick, events, event_id={'left': 769, 'right': 770}, 
                            tmin=0.0, tmax=9.0, baseline=(0, 0), preload=True, verbose=False)
        
        X = epochs.get_data()  # (trials, 2, time)
        y = np.where(epochs.events[:, -1] == 769, 0, 1) # 0 for left, 1 for right

        # 左/右分组
        X_left = X[y == 0]
        X_right = X[y == 1]
        if X_left.size == 0 or X_right.size == 0:
            print(f'被试{sid}: 左或右样本不足，跳过绘图')
            continue

        sfreq = int(raw.info['sfreq'])
        t = np.linspace(0, 9.0, X.shape[2], endpoint=True)

        # 计算平均能量包络
        def avg_env_trials(X):
            c3 = np.stack([band_envelope(tr[0], sfreq=sfreq) for tr in X], axis=0)
            c4 = np.stack([band_envelope(tr[1], sfreq=sfreq) for tr in X], axis=0)
            # 返回：平均值 (Mean), 标准误 (Standard Error)
            return (
                c3.mean(axis=0), c3.std(axis=0) / np.sqrt(c3.shape[0] + 1e-9),
                c4.mean(axis=0), c4.std(axis=0) / np.sqrt(c4.shape[0] + 1e-9),
            )

        l_c3_m, l_c3_se, l_c4_m, l_c4_se = avg_env_trials(X_left)
        r_c3_m, r_c3_se, r_c4_m, r_c4_se = avg_env_trials(X_right)

        # 绘图
        fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
        
        # 左手条件: C3 & C4
        axes[0].set_title(f'被试 {sid} - 左手想象 (C3/C4 Mu 节律 {LOW_CUTOFF}-{HIGH_CUTOFF}Hz)')
        axes[0].plot(t, l_c3_m, label='C3', color='#1f77b4')
        axes[0].fill_between(t, l_c3_m - l_c3_se, l_c3_m + l_c3_se, color='#1f77b4', alpha=0.2)
        axes[0].plot(t, l_c4_m, label='C4', color='#2ca02c')
        axes[0].fill_between(t, l_c4_m - l_c4_se, l_c4_m + l_c4_se, color='#2ca02c', alpha=0.2)
        axes[0].set_ylabel('能量包络(归一化)')
        axes[0].legend(loc='upper right')
        
        # 右手条件: C3 & C4
        axes[1].set_title(f'被试 {sid} - 右手想象 (C3/C4 Mu 节律 {LOW_CUTOFF}-{HIGH_CUTOFF}Hz)')
        axes[1].plot(t, r_c3_m, label='C3', color='#1f77b4')
        axes[1].fill_between(t, r_c3_m - r_c3_se, r_c3_m + r_c3_se, color='#1f77b4', alpha=0.2)
        axes[1].plot(t, r_c4_m, label='C4', color='#2ca02c')
        axes[1].fill_between(t, r_c4_m - r_c4_se, r_c4_m + r_c4_se, color='#2ca02c', alpha=0.2)
        axes[1].set_ylabel('能量包络(归一化)')
        axes[1].set_xlabel('时间 (秒)')
        axes[1].legend(loc='upper right')
        
        for ax in axes:
            # 突出显示运动想象主要发生的 4.0 到 5.0 秒窗口
            ax.axvspan(4.0, 5.0, color='gray', alpha=0.12, label='想象期4-5s')
            ax.grid(True, alpha=0.3)
        
        fig.suptitle(f'左右手想象 Mu 节律响应：C3/C4（{LOW_CUTOFF}-{HIGH_CUTOFF}Hz, 0-9秒，训练数据）')
        fig.tight_layout(rect=[0, 0, 1, 0.96])
        out_path = os.path.join(out_dir, f'subject_{sid}_C3C4_Mu_response_train.png')
        fig.savefig(out_path, dpi=150)
        plt.close(fig)
        print(f'保存: {out_path}')


if __name__ == '__main__':
    # ----------------------------------------------------------------
    # **** 根据您的路径设置：绝对路径 D:\code\Python\library_BCI\MI\CSP+SVM\数据集2a\2a_gdf ****
    # GDF 文件直接在 base_dir 下，无 train/test 子文件夹
    # 标签从 GDF 事件中提取，无需外部 CSV 文件
    # ----------------------------------------------------------------
    
    # 绝对路径设置
    base = r'D:\code\Python\library_BCI\MI\CSP+SVM\数据集2a\2a_gdf'
    
    # 绘图输出目录（相对 base 或绝对，根据需要调整）
    output = os.path.join(os.path.dirname(base), 'plots_c3c4')
    
    try:
        plot_subject_c3c4(base, output)
    except FileNotFoundError as e:
        print(f"\n致命错误: {e}")
        print("请检查您的文件路径设置以及 GDF 文件是否存在。")